# Coppock and Green (2015): Is Voting Habit Forming? New Evidence from Experiments and Regression Discontinuities
# Helper functions


bias_finder <- function(net_migrants, cace=.2, N_c = 500){
  itt = cace*N_c
  ittd <- N_c - net_migrants
  biased <- itt/ittd
  return(biased - cace)
}


complier.control.mean <- function(Y,D,Z){
  pr.c <- mean(D[Z==1]) - mean(D[Z==0])
  pr.at <- mean(D[Z==0])
  pr.nt <- 1 - pr.c - pr.at
  Y0.at <- mean(Y[D==1 & Z==0])
  Y0.nt <- mean(Y[D==0 & Z==1])
  Y0 <- mean(Y[Z==0])
  Y0.c <- (Y0 - pr.at*Y0.at - pr.nt*Y0.nt)/pr.c
  return(Y0.c)
}


cl   <- function(dat,fm, cluster){
  require(sandwich, quietly = TRUE)
  require(lmtest, quietly = TRUE)
  M <- length(unique(cluster))
  N <- length(cluster)
  K <- fm$rank
  dfc <- (M/(M-1))*((N-1)/(N-K))
  uj  <- apply(estfun(fm),2, function(x) tapply(x, cluster, sum));
  vcovCL <- dfc*sandwich(fm, meat=crossprod(uj)/N)
  coeftest(fm, vcovCL) }


state.cace.estimator <-function(data, upstream, downstream, bandwidth, state.abbrev=NULL){
  require(AER)
  require(lmtest)
  if(!is.null(state.abbrev)){ data <- subset(data, state %in% state.abbrev)}
  days <- data[,paste0("days_", upstream)]
  treat <- data[,paste0("treat_", upstream)]
  voted_up <- data[,paste0("voted", upstream)]
  voted_down <- data[,paste0("voted", downstream)]
  voted_down_lag <- data[,paste0("voted", downstream, "_lag")]
  if((sum(voted_up) <500) | (sum(voted_down) <500)){return("Error: not enough data")}
  fit.frame <- subset(data.frame(days, treat, voted_up, voted_down, voted_down_lag), abs(days) <bandwidth)
  fit.2sls <- ivreg(voted_down ~ voted_up + days + voted_up*days + voted_down_lag 
                    | treat + days + treat*days+ voted_down_lag, data=fit.frame)
  
  cace <- fit.2sls$coefficients[2]
  se <- coeftest(fit.2sls, vcov=vcovHC)[2,2]
  
  return(list(summary=summary(fit.2sls), cace=cace, se=se))
}

state.cace.estimator.order <-function(data, upstream_treat,upstream_days, upstream_voted, downstream_voted, downstream_lag, order=1, lag="no",bandwidth, state.abbrev=NULL){
  require(AER)
  require(lmtest)
  if(!is.null(state.abbrev)){ data <- subset(data, state %in% state.abbrev)}

  ##
  days <- data[,upstream_days]
  treat <- data[,upstream_treat]
  voted_up <- data[,upstream_voted]
  voted_down <- data[,downstream_voted]
  voted_down_lag <- data[,downstream_lag]
  
  ##
 # voted_down_lag <- data[,paste0("voted", downstream, "_lag")]
  fit.frame <- subset(data.frame(days, treat, voted_up, voted_down, voted_down_lag), abs(days) <bandwidth)
  fit.frame <- within(fit.frame,{
    days_2 <- days*days
    days_3 <- days*days*days
   # days_4 <- days*days*days*days
  })
  
  if(order==0 & lag=="no"){
    fit.2sls <- ivreg(voted_down ~ voted_up 
                      | treat, data=fit.frame)    
  }

  if(order==1 & lag=="no"){
    fit.2sls <- ivreg(voted_down ~ voted_up + days + voted_up*days 
                      | treat + days + treat*days, data=fit.frame)
  }  
  
  if(order==2 & lag=="no"){
    fit.2sls <- ivreg(voted_down ~ voted_up + days + voted_up*days + days_2 + voted_up*days_2 
                      | treat + days + treat*days + days_2 + treat*days_2, data=fit.frame)
  }
  if(order==3 & lag=="no"){
    fit.2sls <- ivreg(voted_down ~ voted_up + days + voted_up*days + days_2 + voted_up*days_2 + days_3 + voted_up*days_3 
                      | treat + days + treat*days + days_2 + treat*days_2 + days_3 + treat*days_3, data=fit.frame)
  }
  
  if(order==0 & lag=="yes"){
    fit.2sls <- ivreg(voted_down ~ voted_up + voted_down_lag
                      | treat+ voted_down_lag, data=fit.frame)    
  }
  
  if(order==1 & lag=="yes"){
    fit.2sls <- ivreg(voted_down ~ voted_up + days + voted_up*days + voted_down_lag
                      | treat + days + treat*days + voted_down_lag, data=fit.frame)
  }  
  
  if(order==2 & lag=="yes"){
    fit.2sls <- ivreg(voted_down ~ voted_up + days + voted_up*days + days_2 + voted_up*days_2 + voted_down_lag
                      | treat + days + treat*days + days_2 + treat*days_2+ voted_down_lag, data=fit.frame)
  }
  if(order==3 & lag=="yes"){
    fit.2sls <- ivreg(voted_down ~ voted_up + days + voted_up*days + days_2 + voted_up*days_2 + days_3 + voted_up*days_3 + voted_down_lag
                      | treat + days + treat*days + days_2 + treat*days_2 + days_3 + treat*days_3 + voted_down_lag, data=fit.frame)
  }
  
  
  cace <- fit.2sls$coefficients[2]
  se <- coeftest(fit.2sls, vcov=vcovHC)[2,2]
  
  return(list(summary=try(summary(fit.2sls),silent=T), cace=cace, se=se))
}


state.cace.estimator.primary <-function(data, upstream_days, upstream_treat, upstream_voted, 
                                        downstream_voted, downstream_lag, bandwidth, state.abbrev=NULL){
  require(AER)
  require(lmtest)
  if(!is.null(state.abbrev)){ data <- subset(data, state %in% state.abbrev)}
  days <- data[,upstream_days]
  treat <- data[,upstream_treat]
  voted_up <- data[,upstream_voted]
  voted_down <- data[,downstream_voted]
  voted_down_lag <- data[,downstream_lag]
  if((sum(voted_up, na.rm = T) <500) | (sum(voted_down, na.rm=T) <500)){return(
    list(summary=list(coefficients = matrix(NA, 4, 4)), cace=NA, se=NA, summary.itt=list(coefficients = matrix(NA, 4, 4)), coef.itt = NA, se.itt=NA )
  )}
  fit.frame <- subset(data.frame(days, treat, voted_up, voted_down, voted_down_lag), abs(days) <bandwidth)
  fit.2sls <- ivreg(voted_down ~ voted_up + days + voted_up*days + voted_down_lag 
                    | treat + days + treat*days+ voted_down_lag, data=fit.frame)
  itt.mod = lm(voted_down ~treat + days + treat*days+ voted_down_lag, data=fit.frame)
  cace <- fit.2sls$coefficients[2]
  se <- coeftest(fit.2sls, vcov=vcovHC)[2,2]
  itt.coef = itt.mod$coefficients[2]
  se.itt <- coeftest(itt.mod, vcov=vcovHC)[2,2]
  
  return(list(summary=summary(fit.2sls), cace=cace, se=se, summary.itt=summary(itt.mod), coef.itt = itt.coef, se.itt=se.itt ))
}


meta.maker <- function(x, y){
  x <- x[!is.na(x)]
  y <- y[!is.na(y)]
  meta.results <- meta.summaries(x,y)
  return(list(est=meta.results$summary, se=meta.results$se.summary))
}


commarobust <- function(fit){
  require(lmtest)
  require(sandwich)
  coeftest(fit,vcovHC(fit))
}

getrobustses <- function(fit){
  robustfit <- commarobust(fit)
  return(robustfit[,2])
}

makerobustseslist <- function(fitlist){
  return(lapply(fitlist, FUN=getrobustses) )
}



